package org.teamh.gui.engine.model;

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.zip.DataFormatException;
import java.util.zip.Inflater;

import org.lwjgl.BufferUtils;
import org.lwjgl.opengl.GL11;
import org.teamh.gui.engine.model.TextureLoader.RawData;

/**
 * This class supports basic png image loading functions. 
 * @author richard
 *
 */
public class PNGLoader {

	private PNGLoader() {
	}

	/**
	 * Returns a raw data object, which contains any extracted raw data of the png file. 
	 * @param input
	 * @param fileSize
	 * @return
	 */
	public static RawData getPNG(InputStream input, int fileSize) {
		DataInputStream in = new DataInputStream(input);

		long signature = 0;
		try {
			signature = in.readLong();
		} catch (IOException e1) {
			System.err.println("Error while reading png signature");
		}

		if (signature != 0x89504e470d0a1a0aL) {
			System.err.println("This is not a .png file");
			return null;
		}

		Image image = new Image();

		while (true) {
			try {
				// read the size of the chunk
				int length = in.readInt();
				if (length < 0) {
					System.err.println("File is too long");
					break;
				}

				// read the type of the chunk
				byte[] chunkType = new byte[4];
				in.readFully(chunkType);

				// read the data of the chunk
				byte[] chunkData = new byte[length];
				in.read(chunkData);

				// read crc but do nothing
				in.readInt();

				Chunk c = new Chunk(chunkType, chunkData);
				image.add(c);
			} catch (IOException e) {
				// happens when the image file ends
				break;
			}
		}

		image.flush();

		if (image.getCompression() != 0 || image.getFilter() != 0) {
			System.err.println("This image is not compressed correctly");
			return null;
		}

		if (image.getInterlace() != 0) {
			System.err.println("interlace is currently unsupported by png");
			return null;
		}

		int width = image.getWidth();
		int height = image.getHeight();

		int srcPixelFormat = 0;
		int dstPixelFormat = 0;

		switch (image.getColorType()) {
		case 0:
			System.err.println("grayscale is currently unsupported by png");
			return null;
		case 2:
			dstPixelFormat = GL11.GL_RGB;
			if (image.getBitDepth() == 8) {
				srcPixelFormat = GL11.GL_RGB;
			} else if (image.getBitDepth() == 16) {
				srcPixelFormat = GL11.GL_RGB16;
			} else {
				System.err.println("Not a valid RGB pixel depth");
				return null;
			}
			break;
		case 3:
			System.err.println("palette is currently unsupported by png");
			return null;
		case 4:
			System.err
					.println("grayscale with alpha is currently unsupported by png");
			return null;
		case 6:
			dstPixelFormat = GL11.GL_RGBA;
			if (image.getBitDepth() == 8) {
				srcPixelFormat = GL11.GL_RGBA;
			} else if (image.getBitDepth() == 16) {
				srcPixelFormat = GL11.GL_RGBA16;
			} else {
				System.err.println("Not a valid RGB pixel depth");
				return null;
			}
			break;
		default:
			System.err.println("Unsupported color type");
			return null;
		}

		return new RawData(width, height, dstPixelFormat, srcPixelFormat, image
				.getDeflatedData());
	}

	/**
	 * Represents a chunk within a image.
	 * 
	 * @author richard
	 * 
	 */
	private static class Chunk {
		byte[] chunkType;
		byte[] chunkData;

		private Chunk(byte[] chunkType, byte[] chunkData) {
			this.chunkType = chunkType;
			this.chunkData = chunkData;
		}

		public byte[] getData() {
			return chunkData;
		}

		public String getType() {
			try {
				return new String(chunkType, "UTF8");
			} catch (UnsupportedEncodingException e) {
				return "";
			}
		}

		public byte getByte(int pos) {
			return chunkData[pos];
		}

		public int getInt(int pos) {
			int number = 0;

			for (int i = 0; i < 4; i++) {
				number |= (chunkData[pos + i] & 0xff) << ((3 - i) * 8);
			}

			return number;
		}
	}

	/**
	 * This class contains the data of the whole *.png image.
	 * 
	 * @author richard
	 * 
	 */
	private static class Image {
		private Chunk ihdr;
		private ArrayList<Chunk> chunks = new ArrayList<Chunk>();
		private boolean isReading;
		private ByteArrayOutputStream compressedData;

		public Image() {
			isReading = true;
			compressedData = new ByteArrayOutputStream();
		}

		/**
		 * Adds another chunk to the image. 
		 * @param chunk
		 */
		public void add(Chunk chunk) {
			if (chunk.getType().equals("IHDR")) {
				ihdr = chunk;
			} else if (isReading && chunk.getType().equals("IDAT")) {
				try {
					compressedData.write(chunk.getData());
				} catch (IOException e) {
				}
			}
			chunks.add(chunk);
		}

		/**
		 * After all chunks are added, the data can be deflated with this method. 
		 * @return
		 */
		public ByteBuffer getDeflatedData() {
			try {
				int bytesPerPixel = getBytesPerPixel();

				ByteBuffer data = BufferUtils.createByteBuffer(getWidth()
						* getHeight() * bytesPerPixel);
				byte[] filtered = new byte[getWidth() * getHeight()
						* bytesPerPixel];

				Inflater inflater = new Inflater();
				inflater.setInput(compressedData.toByteArray());

				byte[] row = new byte[getWidth() * getBytesPerPixel()];

				for (int y = 0; y < getHeight(); y++) {
					byte[] b = new byte[1];

					inflater.inflate(b);
					inflater.inflate(row);
					
					switch (b[0]) {
					case 0:
						for (int x = 0; x < row.length; x++) {
							filtered[y * row.length + x] = row[x];
						}
						break;
					case 1:
						for (int x = 0; x < bytesPerPixel; x++) {
							filtered[y * row.length + x] = row[x];
						}

						for (int x = bytesPerPixel; x < row.length; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + filtered[y
									* row.length + x - bytesPerPixel]);
						}
						break;
					case 2:
						for (int x = 0; x < row.length; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + filtered[(y - 1)
							                 										* row.length + x]);
						}
						break;
					case 3:
						for (int x = 0; x < bytesPerPixel; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + getAverage(
									(byte) 0, filtered[(y - 1) * row.length
											+ x]));
						}

						for (int x = bytesPerPixel; x < row.length; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + getAverage(
									filtered[(y - 1) * row.length + x],
									filtered[y * row.length + x
											- bytesPerPixel]));
						}
						break;
					case 4:
						for (int x = 0; x < bytesPerPixel; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + getPaethPredictor((byte)0, filtered[(y - 1)
							                 										* row.length + x], (byte)0));
						}

						for (int x = bytesPerPixel; x < row.length; x++) {
							filtered[y * row.length + x] = (byte) (row[x] + getPaethPredictor(
									filtered[y * row.length + x
											- bytesPerPixel],
									filtered[(y - 1) * row.length + x],
									filtered[(y - 1) * row.length + x
											- bytesPerPixel]));
						}
						break;
					default:
						System.err.println("Not a valid filter format");
						return null;
					}
				}

				for (int y = 0; y < getHeight() / 2; y++) {
					for (int x = 0; x < row.length; x++) {
						byte b = filtered[y * row.length + x];
						filtered[y * row.length + x] = filtered[(getHeight()
								- y - 1)
								* row.length + x];
						filtered[(getHeight() - y - 1) * row.length + x] = b;
					}
				}

				data.put(filtered);

				return data;
			} catch (DataFormatException e) {
			}

			return null;
		}
		
		/**
		 * Creates an unsigne byte represented by an integer. 
		 * @param b
		 * @return
		 */
		private int getUnsignedByte(byte b) {
			return b + (b < 0 ? 256 : 0);
		}

		/**
		 * Return the averade of two bytes. 
		 * @param a
		 * @param b
		 * @return
		 */
		private byte getAverage(byte a, byte b) {
			int ia = getUnsignedByte(a);
			int ib = getUnsignedByte(b);

			return (byte) ((ia + ib) / 2);
		}

		/**
		 * Retuns the paeth predictor for png extracting. 
		 * @param a
		 * @param b
		 * @param c
		 * @return
		 */
		private byte getPaethPredictor(byte a, byte b, byte c) {
			int ia = getUnsignedByte(a);
			int ib = getUnsignedByte(b);
			int ic = getUnsignedByte(c);
			
			int p = ia + ib - ic; // initial estimate
			int pa = Math.abs(p - ia);
			int pb = Math.abs(p - ib);
			int pc = Math.abs(p - ic);

			if (pa <= pb && pa <= pc) {
				return a;
			} else if (pb <= pc) {
				return b;
			} else {
				return c;
			}
		}

		/**
		 * This method must be called after all the data has been parsed to this
		 * image.
		 */
		public void flush() {
			try {
				compressedData.flush();
			} catch (IOException e) {
			}

			isReading = false;
		}

		public int getWidth() {
			return ihdr.getInt(0);
		}

		public int getHeight() {
			return ihdr.getInt(4);
		}

		public int getBitDepth() {
			return ihdr.getByte(8);
		}

		public int getColorType() {
			return ihdr.getByte(9);
		}

		public int getCompression() {
			return ihdr.getByte(10);
		}

		public int getFilter() {
			return ihdr.getByte(11);
		}

		public int getInterlace() {
			return ihdr.getByte(12);
		}

		/**
		 * Returns the number of byte per pixel. 
		 * @return
		 */
		public int getBytesPerPixel() {
			if (getColorType() == 2) {
				if (getBitDepth() == 8) {
					return 3;
				} else if (getBitDepth() == 16) {
					return 6;
				}
			} else if (getColorType() == 6) {
				if (getBitDepth() == 8) {
					return 4;
				} else if (getBitDepth() == 16) {
					return 8;
				}
			}

			return 0;
		}
	}
}
